from collections.abc import Mapping
from typing import Any, Dict

from typing_inspect import get_args, get_origin
from zkl_pyutils_registries import UNSET

from zkl_pyutils_serialization.basics import get_context_tag_value
from zkl_pyutils_serialization.convert import convert, registry
from zkl_pyutils_serialization.convertors.common import PyObjectFormat
from zkl_pyutils_serialization.tags import DstFormat, DstTypehint, SrcFormat, SrcType
from .format import JsonValueFormat
from .typing import JsonObject


@registry.register({SrcFormat: PyObjectFormat, DstFormat: JsonValueFormat, SrcType: Mapping, DstTypehint: dict})
@registry.register({SrcFormat: JsonValueFormat, DstFormat: PyObjectFormat, SrcType: Mapping, DstTypehint: dict})
def dict_convertor(value: Mapping) -> JsonObject:
    dst_typehint = get_context_tag_value(DstTypehint)
    key_typehint, val_typehint = parse_dict_type_args(dst_typehint)
    return {
        convert(key, dst_typehint=key_typehint)
        : convert(val, dst_typehint=val_typehint)
        for key, val in value.items()}


@registry.register({SrcFormat: PyObjectFormat, DstFormat: JsonValueFormat, SrcType: dict})
@registry.register({SrcFormat: JsonValueFormat, DstFormat: PyObjectFormat, SrcType: dict})
def dict_convertor_passive(value: JsonObject) -> dict:
    if get_context_tag_value(DstTypehint) not in (UNSET, Any):
        raise NotImplementedError
    return {
        convert(key, dst_typehint=str)
        : convert(val)  # dst_typehint=UNSET
        for key, val in value.items()}


# utils

def parse_dict_type_args(typehint: Any):
    origin = get_origin(typehint) or typehint
    origin = dict if origin is Dict else origin
    if origin not in (dict, Mapping):
        raise NotImplementedError

    args = get_args(typehint, evaluate=True)
    if len(args) == 0:
        key_typehint = UNSET
        val_typehint = UNSET
    elif len(args) == 1:
        key_typehint = args[0]
        val_typehint = UNSET
    elif len(args) == 2:
        key_typehint = args[0]
        val_typehint = args[1]
    else:
        raise TypeError(f'Invalid type hint {typehint} with more than 2 args')
    return key_typehint, val_typehint
